from nodes import DocumentLoaderNode,ChunkDocumentNode,EmbeddingNode,CreateIndexNode,EmbedQueryNode,RetrieveDocumentNode,GenerateAnswerNode

from pocketflow import Flow




def rag_generate_knowledge_base():
    document_loader = DocumentLoaderNode()
    chunk_document = ChunkDocumentNode()
    embedding = EmbeddingNode()
    create_index = CreateIndexNode()
    
    document_loader >> chunk_document >> embedding >> create_index

    flow = Flow(start=document_loader)
    return flow


def rag_generate_answer():
    embed_query = EmbedQueryNode()
    retrieve_document = RetrieveDocumentNode()
    generate_answer = GenerateAnswerNode()

    embed_query >> retrieve_document >> generate_answer

    flow = Flow(start=embed_query)
    return flow

# 初始化两个flow
knowledge_base_flow = rag_generate_knowledge_base()
answer_flow = rag_generate_answer()







